import os
import pandas as pd
import dicom2nifti
from tqdm import tqdm

def ct_or_cta(name):
    '''
        该函数用于判断是ct数据还是cta数据，用于洗数据时候做区分，事实上如果你的任务不需要区分ct和cta，或者有其他的分类需求，修改这里的判断规则
        其他代码部分自己适当调整
    '''
    name = name.lower()
    if 'non contrast' in name or '.nc' in name:
        return 0
    elif 'cta' in name or 'angio' in name or 'dsa' in name:
        return 1
    else:
        return 0

def dcmtonii(dcm_path, id, date, name, modal, nii_path):
    '''
        调用dicom2nifti这个包的dicom_series_to_nifti函数即可，把原始dicom的文件夹（一个患者一个日期一个模态下的一整个序列）以及目标文件名（一个患者一个文件）传入即可
        第三个参数用于控制数据的翻转，选填False就不翻转了
    '''
    dcm_path = os.path.join(dcm_path, id, date, name)
    if modal == 0:
        nii_path = os.path.join(nii_path, 'ct')
    else:
        nii_path = os.path.join(nii_path, 'cta')
    if not os.path.exists(nii_path):
        os.makedirs(nii_path)
    dicom2nifti.dicom_series_to_nifti(dcm_path, os.path.join(nii_path, id+'.nii'), False)
    
if __name__ == '__main__':
    #   DCM为原始数据集路径，NII为目标保存路径，前期需要对所有数据进行记录，保存在need_clip这个excel中
    #   excel不需要表头，第一列为患者病历号（数据名），第二列为一个数字，如果该数据在筛查过程中不符合要求，需要剔除，则记录-1，其余的可以记录为任意≥0的数字
    #   为了方便后面的工作，该表格也可以用于记录数据的有效范围之类的，反正只要关心，不要的数据在第二列写-1就行。
    DCM_PATH = r'xxx'
    NII_PATH = r'xxx'
    if not os.path.exists(NII_PATH):
        os.makedirs(NII_PATH)
        
    NEED_CLIP = r'need_clip.xlsx'

    data_info = pd.read_excel(NEED_CLIP, dtype=str, header=None)
    ok_or_nok = [min(int(s),0) for s in list(data_info[1])]

    patient_list = os.listdir(DCM_PATH)
    ok_dict = {k: v for k, v in zip(patient_list, ok_or_nok)}

    assert patient_list == list(data_info[0]), 'list error'
    #   dicom文件的保存一般是下面这个层次
    #   patient->date->modal->xx.dcm
    for patient_id in tqdm(patient_list):
        if ok_dict[patient_id] is -1:
            continue
        else:
            #   这里我前期已经把数据集筛的，每个患者只剩下一个日期，每个日期里只保留两个模态了，因此有了如下的断言
            date = os.listdir(os.path.join(DCM_PATH, patient_id))
            assert len(date) == 1, '{} not delete clearly, {}, {}'.format(patient_id, date, len(date))
            date = date[0]
            modal_list = os.listdir(os.path.join(DCM_PATH, patient_id, date))
            
            modal_name_1 = modal_list[0]
            modal_name_2 = modal_list[1]
            modal_1 = ct_or_cta(modal_name_1)
            modal_2 = ct_or_cta(modal_name_2)
            assert modal_1 + modal_2 == 1, 'modal repeat or loss'
            #   dcm转换nii的函数
            dcmtonii(DCM_PATH, patient_id, date, modal_name_1, modal_1, NII_PATH)
            dcmtonii(DCM_PATH, patient_id, date, modal_name_2, modal_2, NII_PATH)
